'''
Time Surface Generation from RAW files
Simplified version - with color options
'''
import cv2
import numpy as np
import os
import argparse
from tqdm import tqdm
from metavision_core.event_io import EventsIterator

np.random.seed(0)

def parse_argument():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--raw_file_dir', type=str,
                        default='H:\Datasets/test/raw',
                        help='The path of a raw dir')
    parser.add_argument('--save_dir', type=str,
                        default='H:\Datasets/test/timesurface',
                        help='Output directory for time surface images')
    parser.add_argument('--time_interval_us', type=int, default=50000,
                        help='Time interval in microseconds')
    parser.add_argument('--color_mode', type=str, default='polarization',
                        choices=['grayscale', 'jet', 'hot', 'polarization'],
                        help='Color mode for time surface visualization')
    return parser

def get_reader(file_path, time_interval_us):
    assert os.path.exists(file_path), 'The file \'{}\' does not exist'.format(file_path)
    mv_iterator = EventsIterator(input_path=file_path, delta_t=time_interval_us)
    return mv_iterator

def create_time_surface(events, height, width, color_mode='grayscale'):
    """
    Create time surface representation from events
    
    Args:
        events: event data with 't', 'x', 'y', 'p' fields
        height, width: image dimensions
        color_mode: visualization mode
    
    Returns:
        time_surface: time surface image
    """
    if color_mode == 'polarization':
        return create_polarization_time_surface(events, height, width)
    else:
        return create_colored_time_surface(events, height, width, color_mode)

def create_colored_time_surface(events, height, width, color_mode):
    """Create colored time surface using colormap"""
    # Initialize time surface with zeros
    time_surface = np.zeros((height, width), dtype=np.float64)
    
    if len(events) == 0:
        if color_mode == 'grayscale':
            return (time_surface * 255).astype(np.uint8)
        else:
            return np.zeros((height, width, 3), dtype=np.uint8)
    
    # Filter valid events
    valid_mask = (events['x'] >= 0) & (events['x'] < width) & \
                 (events['y'] >= 0) & (events['y'] < height)
    valid_events = events[valid_mask]
    
    if len(valid_events) == 0:
        if color_mode == 'grayscale':
            return (time_surface * 255).astype(np.uint8)
        else:
            return np.zeros((height, width, 3), dtype=np.uint8)
    
    # For each pixel, keep only the latest event time
    for i in range(len(valid_events)):
        x, y, t = valid_events['x'][i], valid_events['y'][i], valid_events['t'][i]
        if time_surface[y, x] < t:
            time_surface[y, x] = t
    
    # Linear normalization: map time range to [0, 1]
    non_zero_mask = time_surface > 0
    if np.any(non_zero_mask):
        min_time = np.min(time_surface[non_zero_mask])
        max_time = np.max(time_surface[non_zero_mask])
        if max_time > min_time:
            time_surface[non_zero_mask] = (time_surface[non_zero_mask] - min_time) / (max_time - min_time)
        else:
            time_surface[non_zero_mask] = 1.0
    
    # Apply colormap
    if color_mode == 'grayscale':
        return (time_surface * 255).astype(np.uint8)
    else:
        # Convert to uint8 first
        gray_surface = (time_surface * 255).astype(np.uint8)
        
        # Apply OpenCV colormap
        if color_mode == 'jet':
            colored_surface = cv2.applyColorMap(gray_surface, cv2.COLORMAP_JET)
        elif color_mode == 'hot':
            colored_surface = cv2.applyColorMap(gray_surface, cv2.COLORMAP_HOT)
        
        # Set background (no events) to black
        mask = gray_surface == 0
        colored_surface[mask] = [0, 0, 0]
        
        return colored_surface

def create_polarization_time_surface(events, height, width):
    """Create polarization-based colored time surface"""
    # Initialize surfaces for positive and negative events
    pos_surface = np.zeros((height, width), dtype=np.float64)
    neg_surface = np.zeros((height, width), dtype=np.float64)
    
    if len(events) == 0 or 'p' not in events.dtype.names:
        return np.zeros((height, width, 3), dtype=np.uint8)
    
    # Filter valid events
    valid_mask = (events['x'] >= 0) & (events['x'] < width) & \
                 (events['y'] >= 0) & (events['y'] < height)
    valid_events = events[valid_mask]
    
    if len(valid_events) == 0:
        return np.zeros((height, width, 3), dtype=np.uint8)
    
    # Separate positive and negative events
    pos_events = valid_events[valid_events['p'] == 1]
    neg_events = valid_events[valid_events['p'] == 0]
    
    # Process positive events (red channel)
    for i in range(len(pos_events)):
        x, y, t = pos_events['x'][i], pos_events['y'][i], pos_events['t'][i]
        if pos_surface[y, x] < t:
            pos_surface[y, x] = t
    
    # Process negative events (blue channel)
    for i in range(len(neg_events)):
        x, y, t = neg_events['x'][i], neg_events['y'][i], neg_events['t'][i]
        if neg_surface[y, x] < t:
            neg_surface[y, x] = t
    
    # Normalize each channel
    def normalize_channel(surface):
        non_zero_mask = surface > 0
        if np.any(non_zero_mask):
            min_time = np.min(surface[non_zero_mask])
            max_time = np.max(surface[non_zero_mask])
            if max_time > min_time:
                surface[non_zero_mask] = (surface[non_zero_mask] - min_time) / (max_time - min_time)
            else:
                surface[non_zero_mask] = 1.0
        return surface
    
    pos_surface = normalize_channel(pos_surface)
    neg_surface = normalize_channel(neg_surface)
    
    # Create RGB image
    rgb_image = np.zeros((height, width, 3), dtype=np.uint8)
    rgb_image[:, :, 2] = (pos_surface * 255).astype(np.uint8)  # Red channel for positive
    rgb_image[:, :, 0] = (neg_surface * 255).astype(np.uint8)  # Blue channel for negative
    rgb_image[:, :, 1] = 0  # Green channel stays zero
    
    return rgb_image

if __name__ == '__main__':
    # Get params
    args, _ = parse_argument().parse_known_args(None)
    print(args)

    # Find all raw files
    raw_path_list = []
    for root, dirs, files in os.walk(args.raw_file_dir):
        for file in files:
            if file.endswith('.raw'):
                raw_path_list.append(os.path.join(root, file))
    
    print(f"Found {len(raw_path_list)} raw files")
    print(f"Color mode: {args.color_mode}")

    for raw_path in tqdm(raw_path_list):
        assert os.path.exists(raw_path)
        base_name = os.path.basename(raw_path).split('.')[0]
        print(f"Processing: {base_name}")

        # Create output directory
        save_path = os.path.join(args.save_dir, f"{base_name}_{args.color_mode}")
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        reader = get_reader(raw_path, args.time_interval_us)
        height, width = reader.get_size()
        
        print(f"Image dimensions: {width}x{height}")

        frame_cnt = 0
        for evs in reader:
            if len(evs) == 0:
                continue
                
            min_t, max_t = np.min(evs['t']), np.max(evs['t'])
            
            # Create time surface representation
            time_surface = create_time_surface(evs, height, width, args.color_mode)
            
            # Save image
            filename = f'{save_path}/{frame_cnt:06d}_{min_t}_{max_t}.png'
            cv2.imwrite(filename, time_surface)
            
            frame_cnt += 1
            
            if frame_cnt % 100 == 0:
                print(f"Processed {frame_cnt} frames")
        
        print(f"Completed {base_name}: {frame_cnt} frames generated")

    print("Time surface generation completed!")